#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ChromaDB 路径配置诊断脚本
用于检查 ChromaDB 路径配置问题和数据存储情况
"""

import asyncio
import os
import sys
import yaml
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root / 'src'))

from database.manager import DatabaseManager
from database.chromadb_client import ChromaDBClient

def load_config():
    """加载配置文件"""
    config_path = project_root / 'config' / 'database.yaml'
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

async def diagnose_chromadb_path():
    """诊断 ChromaDB 路径配置问题"""
    print("=== ChromaDB 路径配置诊断 ===")
    
    # 1. 检查配置文件中的路径设置
    config = load_config()
    chromadb_config = config.get('databases', {}).get('chromadb', {})
    configured_path = chromadb_config.get('persist_directory', './data/chroma_db')
    print(f"\n1. 配置文件中的路径: {configured_path}")
    
    # 2. 检查实际存在的目录
    print("\n2. 检查实际存在的目录:")
    data_dir = project_root / 'data'
    if data_dir.exists():
        print(f"   ✅ {data_dir} 存在")
        for subdir in data_dir.iterdir():
            if subdir.is_dir():
                count = len(list(subdir.iterdir()))
                print(f"   - {subdir.name}/: {count} 个子目录")
    else:
        print(f"   ❌ {data_dir} 不存在")
    
    # 3. 检查根目录下的 chroma_db
    root_chroma = project_root / 'chroma_db'
    if root_chroma.exists():
        count = len(list(root_chroma.iterdir()))
        print(f"   ✅ {root_chroma} 存在: {count} 个子目录")
    else:
        print(f"   ❌ {root_chroma} 不存在")
    
    # 4. 初始化 ChromaDB 客户端并检查实际使用的路径
    print("\n3. 初始化 ChromaDB 客户端...")
    try:
        chromadb_client = ChromaDBClient(config)
        await chromadb_client.connect()
        
        # 检查客户端配置
        actual_path = chromadb_client.chromadb_config.get('persist_directory', './data/chroma_db')
        print(f"   客户端实际使用的路径: {actual_path}")
        
        # 检查集合统计
        print("\n4. 检查集合统计:")
        collections = ['memories', 'documents', 'conversations', 'knowledge']
        for collection_name in collections:
            try:
                stats = await chromadb_client.get_collection_stats(collection_name)
                count = stats.get('count', 0)
                print(f"   - {collection_name}: {count} 个文档")
            except Exception as e:
                print(f"   - {collection_name}: 错误 - {e}")
        
        # 5. 测试对话搜索
        print("\n5. 测试对话搜索:")
        try:
            search_results = await chromadb_client.search_conversations(
                query="测试",
                user_id="test_user",
                limit=5
            )
            print(f"   搜索结果数量: {len(search_results)}")
            if search_results:
                print(f"   示例结果: {search_results[0].get('content', '')[:50]}...")
        except Exception as e:
            print(f"   搜索失败: {e}")
        
        await chromadb_client.close()
        
    except Exception as e:
        print(f"   ❌ ChromaDB 客户端初始化失败: {e}")
        import traceback
        traceback.print_exc()
    
    # 6. 检查 Neo4j 连接状态
    print("\n6. 检查 Neo4j 连接状态:")
    try:
        db_manager = DatabaseManager(config)
        await db_manager.initialize()
        
        # 检查 Neo4j 客户端类型
        neo4j_client = db_manager.neo4j_client
        client_type = type(neo4j_client).__name__
        print(f"   Neo4j 客户端类型: {client_type}")
        
        if hasattr(neo4j_client, 'find_related_concepts'):
            # 测试 find_related_concepts
            try:
                concepts = await neo4j_client.find_related_concepts("测试概念")
                print(f"   find_related_concepts 结果: {len(concepts)} 个概念")
            except Exception as e:
                print(f"   find_related_concepts 失败: {e}")
        else:
            print(f"   ❌ Neo4j 客户端没有 find_related_concepts 方法")
        
        await db_manager.close()
        
    except Exception as e:
        print(f"   ❌ Neo4j 连接检查失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(diagnose_chromadb_path())